#!/usr/bin/python3
# -*- coding: utf-8 -*-
import os
import re
import struct

out_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin'
bbox_count_list_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin/bbox_count_list.txt'
bbox_list_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin/bbox_list.txt'
select_list_path = r'/rootfs/media/kasim/DataSet/FaceVideo/CropImageMargin/FilterSelectImage.txt'

select_label = 4


def main():
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    os.system('chmod a+wr {}'.format(out_path))

    file_names = []
    with open(select_list_path, 'r') as file:
        for line in file.readlines():
            file_name = line.strip()
            file_names.append(file_name)

    file_names.sort()

    pattern = re.compile(r'F([-\d]+)B([-\d]+)')

    dir_names = []
    dir_frames_map = {}
    for file_name in file_names:
        dir_name = os.path.dirname(file_name)
        frame_info = os.path.basename(file_name)
        fr_num, b_num = map(int, pattern.search(frame_info).groups())
        frames = dir_frames_map.get(dir_name, None)
        if frames is None:
            dir_names.append(dir_name)
            frames = []
            dir_frames_map[dir_name] = frames
        frames.append(fr_num)

    out_file_bbox_path = os.path.join(out_path, 'select_file_bbox.txt')

    out_list_file_path = os.path.join(out_path, 'select_bbox_list.txt')
    out_video_bbox_file_path = os.path.join(out_path, 'select_bbox.dat')

    out_file_bbox = open(out_file_bbox_path, 'w')

    out_list_file = open(out_list_file_path, 'w')
    out_video_bbox_file = open(out_video_bbox_file_path, 'wb')

    file_bbox_map = {}
    dir_frame_bbox_map = {}
    with open(bbox_count_list_path, 'r') as bbox_count_file:
        with open(bbox_list_path, 'r') as bbox_file:
            for line in bbox_count_file.readlines():
                lines = line.strip().split(',')
                file_dir = os.path.splitext(lines[0])[0]

                frame_bbox_map = dir_frame_bbox_map.get(file_dir, None)
                if frame_bbox_map is None:
                    frame_bbox_map = {}
                    dir_frame_bbox_map[file_dir] = frame_bbox_map
                bbox_count = int(lines[1])
                for _ in range(bbox_count):
                    bbox_infos = bbox_file.readline()
                    bbox_infos = bbox_infos.strip().split()
                    file_name = bbox_infos[0]

                    fr_num, b_num = map(int, pattern.search(file_name).groups())

                    file_path = os.path.join(file_dir, file_name)
                    bbox_info1 = bbox_infos[1]
                    bbox_info2 = bbox_infos[2]

                    bboxes = frame_bbox_map.get(fr_num, None)
                    if bboxes is None:
                        bboxes = [[], []]
                        frame_bbox_map[fr_num] = bboxes

                    bboxes[0].append(bbox_info1)
                    bboxes[1].append(bbox_info2)
                    file_bbox_map[file_path] = [bbox_info1, bbox_info2]

    for file_name in file_names:
        bboxes = file_bbox_map.get(file_name, None)
        if bboxes is None:
            continue
        out_info_str = '{} {} {}\n'.format(file_name, bboxes[0], bboxes[1])
        out_file_bbox.write(out_info_str)

    # for dir_name, frames in dir_frames_map.items():
    for dir_name in dir_names:
        frames = dir_frames_map[dir_name]

        frames.sort()
        frame_bbox_map = dir_frame_bbox_map.get(dir_name, None)
        if frame_bbox_map is None:
            continue

        bbox_count = 0
        for frame in frames:
            bboxes = frame_bbox_map.get(frame, None)
            if bboxes is None:
                continue
            # bboxes0 = bboxes[0]
            bboxes1 = bboxes[1]

            bbox_count += len(bboxes1)
            for bbox in bboxes1:
                bbox = bbox.split(',')
                bbox = [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3]), float(bbox[4])]
                dat = struct.pack('6i1f', frame, select_label, bbox[0], bbox[1], bbox[2], bbox[3], bbox[4])
                # dat_list = struct.unpack('6i1f', dat)
                # print(len(dat), dat_list)
                out_video_bbox_file.write(dat)

        if bbox_count > 0:
            out_info_str = '{}.avi,{}\n'.format(dir_name, bbox_count)
            out_list_file.write(out_info_str)

    out_file_bbox.close()
    out_list_file.close()
    out_video_bbox_file.close()
    os.system('chmod a+wr {}'.format(out_file_bbox_path))
    os.system('chmod a+wr {}'.format(out_list_file_path))
    os.system('chmod a+wr {}'.format(out_video_bbox_file_path))

    print('finish!!!!')


if __name__ == '__main__':
    main()
